import base64
import datetime
import hmac
import random
import time
from dataclasses import dataclass, field
from enum import Enum
from hashlib import sha256
from typing import Any, Dict, Optional


# picc中台请求报文格式
class SessionStatus(Enum):
    ONCE = "once"
    BEGIN = "begin"
    CONTINUE = "continue"
    END = "end"

    @classmethod
    def value_of(cls, value: str):
        for mode in cls:
            if mode.value == value:
                return mode
        raise ValueError(f"invalid session status type value {value}")


class MessageRole(Enum):
    USER = "user"
    SYSTEM = "system"
    ASSISTANT = "assistant"

    @classmethod
    def value_of(cls, value: str):
        for mode in cls:
            if mode.value == value:
                return mode
        raise ValueError(f"invalid messsage role type value {value}")


class PiccModel(Enum):
    QWEN = "qwen"
    DEEPSEEK = "deepseek"

    @classmethod
    def value_of(cls, value: str):
        """
        Get value of given mode.

        :param value: mode value
        :return: mode
        """
        for mode in cls:
            if mode.value == value:
                return mode
        raise ValueError(f"invalid Picc Model type value {value}")


@dataclass
class Session:
    sid: str = ""
    following: str = ""
    status: SessionStatus = SessionStatus.ONCE

    def to_dict(self):
        return {
            "sid": self.sid,
            "following": self.following,
            "status": self.status.value
        }


@dataclass
class ProductParam:
    product_id: str = ""
    scene_id: str = ""
    user_id: str = ""
    request_id: str = ""
    extend_product_params: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self):
        return {
            "product_id": self.product_id,
            "scene_id": self.scene_id,
            "user_id": self.user_id,
            "request_id": self.request_id,
            "extend_product_params": self.extend_product_params
        }


@dataclass
class ModelParams:
    pass


@dataclass
class Params:
    product_params: ProductParam
    model_params: ModelParams

    def to_dict(self):
        return {
            "product_params": self.product_params.to_dict(),
            "model_params": self.model_params.to_dict()
        }


@dataclass
class Message:
    role: MessageRole
    content: str

    def to_dict(self):
        return {
            "role": self.role.value,
            "content": self.content
        }


@dataclass
class Payloads:
    message: list[Message] = field(default_factory=list)

    def to_dict(self):
        return {
            "message": [msg.to_dict() for msg in self.message],
        }


@dataclass
class PayloadsWithTools(Payloads):
    tools: list[str] = field(default_factory=list)

    def to_dict(self):
        dic = super().to_dict()
        dic["tools"] = [tool for tool in self.tools]
        return dic


@dataclass
class PiccRequest:
    session: Session = field(default_factory=Session)
    params: Params = field(default_factory=Params)
    payloads: Payloads = field(default_factory=Payloads)

    def to_dict(self):
        return {
            "session": self.session.to_dict(),
            "params": self.params.to_dict(),
            "payloads": self.payloads.to_dict()
        }


# 针对不同类型的模型，Params.model_params不同，需定制化开发
# 阿里千问
@dataclass
class ModelParamsQWen(ModelParams):
# 选传参数
    use_tools: bool = False
    eos_token_id: Optional[int] = None
    bad_words_ids: list[list[int]] = field(default_factory=list)
    prefix_allowed_tokens_fn: str = "test_function"
    extend_model_params: Dict[str, Any] = field(default_factory=dict)
# 固定参数
    model_id: str = "001_int4"  # qwen模型名称固定为 001_int4
    num_beams: int = 1  # 固定为 1
    num_beam_groups: int = 1  # 固定为 1
    do_sample: bool = True  # 设置为 True 才能设置 top_p
# dify模型配置
    top_p: float = 0.8
    top_k: int = 10
    temperature: float = 0.3
    repetition_penalty: float = 1.0
    diversity_penalty: float = 0.0
    length_penalty: float = 0.0
    max_new_tokens: int = 512
    max_length: int = 512
    num_return_sequences: int = 1

    def to_dict(self):
        return {
            "use_tools": self.use_tools,
            "eos_token_id": self.eos_token_id,
            "bad_words_ids": self.bad_words_ids,
            "prefix_allowed_tokens_fn": self.prefix_allowed_tokens_fn,
            "extend_model_params": self.extend_model_params,
            "model_id": self.model_id,
            "num_beams": self.num_beams,
            "num_beam_groups": self.num_beam_groups,
            "do_sample": self.do_sample,
            "top_p": self.top_p,
            "top_k": self.top_k,
            "temperature": self.temperature,
            "repetition_penalty": self.repetition_penalty,
            "diversity_penalty": self.diversity_penalty,
            "length_penalty": self.length_penalty,
            "max_new_tokens": self.max_new_tokens,
            "max_length": self.max_length,
            "num_return_sequences": self.num_return_sequences,
        }


# Deepseek
@dataclass
class ModelParamsDeepSeek(ModelParams):
# 必传参数4
    model_id: str
    model_secret: str
    timeout: int = 600
# dify模型配置
    top_p: float = 0.96
    top_k: int = 20
    temperature: float = 0.6
    repetition_penalty: float = 1
    max_tokens: int = 10240

    def to_dict(self):
        return {
            "model_id": self.model_id,
            "model_secret": self.model_secret,
            "timeout": self.timeout,
            "top_p": self.top_p,
            "top_k": self.top_k,
            "temperature": self.temperature,
            "repetition_penalty": self.repetition_penalty,
            "max_tokens": self.max_tokens,
        }


# 请求头
def get_request_headers(app_id: str,
                        secret_id: str,
                        secret_key: str,
                        ability_id: str,
                        action: str = 'default',
                        scene_id: str = 'default',
                        http_method: str = 'POST'):
    random_str = str(random.randint(1, 99999999999))
    request_headers = {
        'X-IFAC-Ability': ability_id,
        'X-IFAC-Action': action,
        'X-IFAC-App-Id': app_id,
        'X-IFAC-Nonce': random_str,
        'X-IFAC-Scene-Id': scene_id,
        'X-IFAC-Timestamp': str(int(time.mktime(datetime.datetime.now().timetuple()))),
    }
    request_headers['Authorization'] = get_authorization_header(http_method, request_headers, secret_id, secret_key)
    return request_headers


# 请求头加密字段
def get_authorization_header(http_method: str,
                             canonicalized_headers: dict,
                             secret_id: str,
                             secret_key: str):
    signature_method = 'HmacSHA256'
    canonicalized_headers_str = ''
    for key in sorted(canonicalized_headers.keys()):
        canonicalized_headers_str += key.lower().strip() + '=' + str(canonicalized_headers[key]).strip() + '\n'
    canonicalized_resource_str = '/'
    string_to_sign = http_method.upper() + '\n' + canonicalized_headers_str + canonicalized_resource_str
    signature = base64.b64encode(
        hmac.new(key=secret_key.encode('utf-8'), msg=string_to_sign.encode('utf-8'), digestmod=sha256).digest()
    )
    signature = str(signature, 'utf-8')
    authorization = str(signature_method) + ':' + str(secret_id) + ':' + str(signature)
    return authorization


# 请求体
def get_request_body(
        prompt_messages: list[Message],
        credentials: dict,
        model_parameters: dict,
        is_init: bool = False,
    ) -> PiccRequest:
    model = PiccModel.value_of(credentials["model_base"].lower())
    if is_init:
        model_parameters = {}
    if model is PiccModel.QWEN:
        modelParams = ModelParamsQWen(
            **model_parameters
        )
        payload = PayloadsWithTools(
            message=prompt_messages
        )
    elif model is PiccModel.DEEPSEEK:
        modelParams = ModelParamsDeepSeek(
            model_id=credentials["model_id"],
            model_secret=credentials["model_secret"],
            **model_parameters
        )
        payload = Payloads(
            message=prompt_messages
        )
    else:
        return None

    productParam = ProductParam(
        product_id="dify"
    )

    params = Params(
        product_params=productParam,
        model_params=modelParams
    )

    session = Session(
        status=SessionStatus.ONCE
    )
    request_body = PiccRequest(
        session=session,
        params=params,
        payloads=payload
    )
    return request_body
